# -*- coding: utf-8 -*-
"""
基于Web界面的银行客户查询助手 - 增强版
"""
import os
import asyncio
from typing import Optional
import dashscope
from qwen_agent.agents import Assistant
from qwen_agent.gui import WebUI
import pandas as pd
from sqlalchemy import create_engine
from qwen_agent.tools.base import BaseTool, register_tool
import matplotlib.pyplot as plt
import io
import base64
import time
import numpy as np

# 解决中文显示问题
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'SimSun', 'Arial Unicode MS']  # 优先使用的中文字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 定义资源文件根目录
ROOT_RESOURCE = os.path.join(os.path.dirname(__file__), 'resource')

# 配置数据库连接
MYSQL_USER = os.getenv('MYSQL_USER', 'root')
MYSQL_PASSWORD = os.getenv('MYSQL_PASSWORD', '')
MYSQL_HOST = 'localhost'
MYSQL_PORT = '3306'
DATABASE_NAME = 'bank'  # 数据库名更改为 bank

def create_connection():
    """创建数据库连接"""
    try:
        engine = create_engine(
            f'mysql+mysqlconnector://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}:{MYSQL_PORT}/{DATABASE_NAME}?charset=utf8mb4',
            connect_args={'connect_timeout': 10}, pool_size=10, max_overflow=20
        )
        return engine
    except Exception as e:
        print(f"数据库连接失败: {str(e)}")
        return None

# ====== 更新助手的 system_prompt，添加开场白 ======
system_prompt = """
您好！欢迎使用银行客户查询助手。

我是您的智能助手，可以帮助您高效查询和分析银行客户数据。以下是我的功能：
- 查询客户基本信息（姓名、年龄、职业等）
- 分析客户的资产分布（存款、理财、基金等）
- 查看客户的交易行为和APP使用习惯
- 生成可视化图表以辅助决策
- 使用逻辑回归模型预测客户未来N个月资产提升至100万+的概率

以下是数据库中的表结构说明：

客户基础信息表结构如下：
CREATE TABLE customer_base (
    customer_id VARCHAR(32) PRIMARY KEY COMMENT '客户ID',
    name VARCHAR(100) COMMENT '客户姓名',
    age INT COMMENT '年龄',
    gender VARCHAR(10) COMMENT '性别',
    occupation VARCHAR(100) COMMENT '职业',
    occupation_type VARCHAR(50) COMMENT '职业类型标签（如：企业高管/互联网从业者/私营业主）',
    monthly_income DECIMAL(12,2) COMMENT '月收入',
    open_account_date VARCHAR(10) COMMENT '开户日期',
    lifecycle_stage VARCHAR(50) COMMENT '客户生命周期',
    marriage_status VARCHAR(20) COMMENT '婚姻状态',
    city_level VARCHAR(20) COMMENT '城市等级（一线/二线城市）',
    branch_name VARCHAR(100) COMMENT '开户网点'
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='客户基础信息表';

客户行为资产表结构如下：
CREATE TABLE customer_behavior_assets (
    id VARCHAR(32) PRIMARY KEY COMMENT '主键ID',
    customer_id VARCHAR(32) COMMENT '客户ID（关联customer_base表）',
    total_assets DECIMAL(16,2) COMMENT '总资产',
    deposit_balance DECIMAL(16,2) COMMENT '存款余额',
    financial_balance DECIMAL(16,2) COMMENT '理财余额',
    fund_balance DECIMAL(16,2) COMMENT '基金余额',
    insurance_balance DECIMAL(16,2) COMMENT '保险余额',
    asset_level VARCHAR(20) COMMENT '资产分层（50万以下、50-80万、80-100万、100万+）',
    deposit_flag TINYINT COMMENT '是否持有存款（1是0否）',
    financial_flag TINYINT COMMENT '是否持有理财（1是0否）',
    fund_flag TINYINT COMMENT '是否持有基金（1是0否）',
    insurance_flag TINYINT COMMENT '是否持有保险（1是0否）',
    product_count INT COMMENT '持有产品数量',
    financial_repurchase_count INT COMMENT '近1年理财复购次数',
    credit_card_monthly_expense DECIMAL(12,2) COMMENT '信用卡月均消费',
    investment_monthly_count INT COMMENT '月均投资交易次数',
    app_login_count INT COMMENT 'APP月均登录次数',
    app_financial_view_time INT COMMENT '理财页面月均停留时长(秒)',
    app_product_compare_count INT COMMENT '产品对比点击次数',
    last_app_login_time VARCHAR(19) NULL COMMENT '最近APP登录时间',
    last_contact_time VARCHAR(19) NULL COMMENT '最近联系时间',
    contact_result VARCHAR(50) COMMENT '联系结果',
    marketing_cool_period VARCHAR(10) COMMENT '营销冷却期（30天）',
    stat_month VARCHAR(7) COMMENT '统计月份（YYYY-MM）',
    FOREIGN KEY (customer_id) REFERENCES customer_base(customer_id),
    UNIQUE KEY uk_customer_month (customer_id, stat_month)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='客户行为资产表';

每当 exc_sql 工具返回 markdown 表格和图片时，你必须原样输出工具返回的全部内容（包括图片），不要只总结表格，也不要省略图片。这样用户才能直接看到表格和图片。

重要提醒：
1. 当工具返回包含图表的结果时，必须在你的回答中完整显示这些图表
2. 在给出分析结论后，如果有相关图表，请再次展示图表以便用户查看
3. 图表是数据分析的重要组成部分，不能省略或仅用文字描述替代
4. 确保图表链接格式正确：![图表标题](图片路径)

知识库中是以前做过的示例，如果有类似的需要进行逻辑上的参考。
"""

functions_desc = [
    {
        "name": "exc_sql",
        "description": "对于生成的SQL，进行SQL查询",
        "parameters": {
            "type": "object",
            "properties": {
                "sql_input": {
                    "type": "string",
                    "description": "生成的SQL语句",
                }
            },
            "required": ["sql_input"],
        },
    },
]

# ====== 会话隔离 DataFrame 存储 ======
# 用于存储每个会话的 DataFrame，避免多用户数据串扰
_last_df_dict = {}

def get_session_id(kwargs):
    """根据 kwargs 获取当前会话的唯一 session_id，这里用 messages 的 id"""
    messages = kwargs.get('messages')
    if messages is not None:
        return id(messages)
    return None

# ====== exc_sql 工具类实现 ======
@register_tool('exc_sql')
class ExcSQLTool(BaseTool):
    """
    SQL查询工具，执行传入的SQL语句并返回结果，并自动进行可视化。
    """
    description = '对于生成的SQL，进行SQL查询，并自动可视化'
    parameters = [{
        'name': 'sql_input',
        'type': 'string',
        'description': '生成的SQL语句',
        'required': True
    }, {
        'name': 'need_visualization',
        'type': 'boolean',
        'description': '是否需要可视化图表，默认为True',
        'required': False
    }]

    def call(self, params: str, **kwargs) -> str:
        import json
        args = json.loads(params)
        sql_input = args['sql_input']
        need_visualization = args.get('need_visualization', True)  # 默认需要可视化
        database = args.get('database', DATABASE_NAME)  # 使用配置的数据库名称
        engine = create_connection()
        
        try:
            df = pd.read_sql(sql_input, engine)
            
            # 检查是否是股票对比查询，如果是则不需要可视化
            if "对比" in sql_input or "比较" in sql_input:
                need_visualization = False
            
            # 生成Markdown表格，显示前5行和后5行
            if len(df) <= 10:
                md = df.to_markdown(index=False)
            else:
                # 显示前5行和后5行
                head_df = df.head(5)
                tail_df = df.tail(5)
                md = head_df.to_markdown(index=False) + "\n...\n" + tail_df.to_markdown(index=False)
            
            # 如果不需要可视化，直接返回表格数据
            if not need_visualization:
                # 如果是股票对比查询，添加涨跌幅计算公式和结果
                if "对比" in sql_input or "比较" in sql_input:
                    pct_chg_formula = "\n\n涨跌幅计算公式：\n涨跌幅(%) = (当日收盘价 - 昨日收盘价) / 昨日收盘价 × 100%\n\n"
                    if 'pct_chg' in df.columns:
                        pct_chg_formula += "查询结果中的涨跌幅(%)字段即为按上述公式计算的结果。\n"
                    return f"{md}{pct_chg_formula}"
                return md
            
            # 添加数据描述信息
            describe_md = "数据统计信息：\n" + df.describe().to_markdown()
            
            # 如果只有一行或没有数据，则不生成图表
            if len(df) <= 1:
                if len(df) > 0:
                    return f"{md}\n\n{describe_md}"
                return md
            
            # 自动创建目录
            save_dir = os.path.join(os.path.dirname(__file__), 'image_show')
            os.makedirs(save_dir, exist_ok=True)
            filename = f'bar_{int(time.time()*1000)}.png'
            save_path = os.path.join(save_dir, filename)
            
            # 生成图表
            generate_chart_png(df, save_path)
            img_path = os.path.join('image_show', filename)
            img_md = f'![柱状图]({img_path})'
            
            return f"{md}\n\n{img_md}"
        except Exception as e:
            return f"查询执行失败: {str(e)}"

def generate_chart_png(df, save_path):
    """生成柱状图并保存"""
    plt.figure(figsize=(10, 6))
    df.iloc[:, -1].plot(kind='bar', color='skyblue')
    plt.title('数据分布')
    plt.xlabel('索引')
    plt.ylabel('值')
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()

# ====== 新增：逻辑回归预测工具 ======
@register_tool('lr_predict')
class LRPredictTool(BaseTool):
    """
    逻辑回归预测工具，预测客户未来N个月资产提升至100万+的概率。
    """
    description = '使用逻辑回归模型预测客户未来N个月资产提升至100万+的概率'
    parameters = [{
        'name': 'customer_id',
        'type': 'string',
        'description': '客户ID',
        'required': True
    }, {
        'name': 'prediction_period',
        'type': 'integer',
        'description': '预测周期（月数），默认为3',
        'required': False
    }]

    def call(self, params: str, **kwargs) -> str:
        import json
        args = json.loads(params)
        customer_id = args['customer_id']
        prediction_period = args.get('prediction_period', 3)  # 默认预测3个月
        
        # 从数据库获取客户数据
        engine = create_connection()
        if engine is None:
            return "数据库连接失败"
        
        # 获取客户历史数据
        query = f"""
        SELECT cb.name, cb.age, cb.gender, cb.monthly_income, cb.occupation, 
               cba.total_assets, cba.deposit_balance, cba.financial_balance, 
               cba.fund_balance, cba.insurance_balance, cba.product_count,
               cba.financial_repurchase_count, cba.credit_card_monthly_expense,
               cba.investment_monthly_count, cba.app_login_count,
               cba.app_financial_view_time, cba.app_product_compare_count,
               cba.stat_month
        FROM customer_base cb
        JOIN customer_behavior_assets cba ON cb.customer_id = cba.customer_id
        WHERE cb.customer_id = '{customer_id}'
        ORDER BY cba.stat_month
        """
        
        try:
            df = pd.read_sql(query, engine)
            if df.empty:
                return f"未找到客户ID为 {customer_id} 的数据"
            
            # 计算增长率
            df['prev_assets'] = df['total_assets'].shift(1)
            df['asset_growth'] = df['total_assets'] - df['prev_assets']
            df['growth_rate'] = df['asset_growth'] / df['prev_assets']
            
            # 计算平均增长率
            avg_growth_rate = df['growth_rate'].mean()
            
            # 计算预测资产
            current_assets = df['total_assets'].iloc[-1]
            predicted_assets = current_assets * (1 + avg_growth_rate) ** prediction_period
            
            # 判断是否达到100万+
            target_reached = int(predicted_assets >= 1000000)
            
            # 特征工程 - 使用01编码处理分类变量
            features_for_clustering = [
                'age', 'monthly_income', 'total_assets', 'deposit_balance', 
                'financial_balance', 'fund_balance', 'insurance_balance',
                'product_count', 'financial_repurchase_count', 
                'credit_card_monthly_expense', 'investment_monthly_count',
                'app_login_count', 'app_financial_view_time', 
                'app_product_compare_count', 'avg_growth_rate'
            ]
            
            # 分类特征
            categorical_features = [
                'gender', 'occupation', 'city_level', 'branch_name', 'asset_level'
            ]
            
            # 处理数值特征
            numeric_data = df[features_for_clustering].fillna(0)
            
            # 处理分类特征 - 使用01编码
            categorical_data = df[categorical_features].copy()
            label_encoders = {}
            
            for feature in categorical_features:
                le = LabelEncoder()
                categorical_data[feature] = le.fit_transform(categorical_data[feature].astype(str))
                label_encoders[feature] = le
            
            # 合并所有特征
            X = pd.concat([numeric_data, categorical_data], axis=1)
            
            # 训练逻辑回归模型
            from sklearn.linear_model import LogisticRegression
            from sklearn.preprocessing import StandardScaler
            
            scaler = StandardScaler()
            X_scaled = scaler.fit_transform(X)
            
            # 创建目标变量
            y = [target_reached] * len(X_scaled)
            
            # 训练模型
            lr = LogisticRegression(max_iter=2000)
            lr.fit(X_scaled, y)
            
            # 获取系数
            coef_df = pd.DataFrame({
                'feature': X.columns,
                'coefficient': lr.coef_[0]
            }).sort_values('coefficient', key=abs, ascending=False)
            
            # 可视化系数
            plt.figure(figsize=(10, 8))
            colors = ['blue' if x > 0 else 'red' for x in coef_df['coefficient']]
            plt.barh(range(len(coef_df)), coef_df['coefficient'], color=colors)
            plt.yticks(range(len(coef_df)), coef_df['feature'])
            plt.xlabel('Coefficient Value')
            plt.title('Logistic Regression Coefficients for Predicting 1M+ Assets in {} Months'.format(prediction_period))
            plt.axvline(x=0, color='black', linewidth=0.5)
            
            # 添加图例
            red_patch = plt.Rectangle((0,0),1,1, color='red')
            blue_patch = plt.Rectangle((0,0),1,1, color='blue')
            plt.legend([blue_patch, red_patch], ['Positive Coefficient', 'Negative Coefficient'])
            
            plt.tight_layout()
            save_dir = os.path.join(os.path.dirname(__file__), 'image_show')
            os.makedirs(save_dir, exist_ok=True)
            filename = f'lr_coefficients_{customer_id}_{prediction_period}.png'
            save_path = os.path.join(save_dir, filename)
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            plt.close()
            
            # 生成结果
            result = f"""
            客户ID: {customer_id}
            客户姓名: {df['name'].iloc[-1]}
            当前总资产: {current_assets:.2f}
            预测{prediction_period}个月后总资产: {predicted_assets:.2f}
            达到100万+概率: {target_reached}
            
            逻辑回归系数:
            {coef_df.to_markdown(index=False)}
            
            系数可视化图表已保存到 {filename}
            """
            
            return result
            
        except Exception as e:
            return f"预测执行失败: {str(e)}"



# ====== 新增：LightGBM 客户价值因素分析工具 ======
@register_tool('lightgbm_analysis')
class LightGBMAnalysisTool(BaseTool):
    """
    使用 LightGBM 模型分析客户未来 3 个月资产提升至 100 万+的概率，并输出特征重要性排序。
    """
    description = '使用 LightGBM 模型分析客户未来 3 个月资产提升至 100 万+的概率，并输出特征重要性'
    parameters = [{
        'name': 'prediction_period',
        'type': 'integer',
        'description': '预测周期（月数），默认为 3',
        'required': False
    }]

    def call(self, params: str, **kwargs) -> str:
        import json
        import lightgbm as lgb
        from sklearn.model_selection import train_test_split
        from sklearn.preprocessing import LabelEncoder, StandardScaler
        import matplotlib.pyplot as plt

        args = json.loads(params)
        prediction_period = args.get('prediction_period', 3)  # 默认预测 3 个月

        # 从数据库获取客户数据
        engine = create_connection()
        if engine is None:
            return "数据库连接失败"

        # 获取客户历史数据
        query = f"""
        SELECT cb.customer_id, cb.age, cb.gender, cb.monthly_income, cb.occupation, 
               cba.total_assets, cba.deposit_balance, cba.financial_balance, 
               cba.fund_balance, cba.insurance_balance, cba.product_count,
               cba.financial_repurchase_count, cba.credit_card_monthly_expense,
               cba.investment_monthly_count, cba.app_login_count,
               cba.app_financial_view_time, cba.app_product_compare_count,
               cba.stat_month
        FROM customer_base cb
        JOIN customer_behavior_assets cba ON cb.customer_id = cba.customer_id
        ORDER BY cb.customer_id, cba.stat_month
        """

        try:
            df = pd.read_sql(query, engine)
            if df.empty:
                return "未找到任何客户数据"

            # 特征工程
            features_for_clustering = [
                'age', 'monthly_income', 'total_assets', 'deposit_balance', 
                'financial_balance', 'fund_balance', 'insurance_balance',
                'product_count', 'financial_repurchase_count', 
                'credit_card_monthly_expense', 'investment_monthly_count',
                'app_login_count', 'app_financial_view_time', 
                'app_product_compare_count'
            ]

            categorical_features = ['gender', 'occupation']

            # 处理数值特征
            numeric_data = df[features_for_clustering].fillna(0)

            # 处理分类特征 - 使用 LabelEncoder 编码
            categorical_data = df[categorical_features].copy()
            label_encoders = {}
            for feature in categorical_features:
                le = LabelEncoder()
                categorical_data[feature] = le.fit_transform(categorical_data[feature].astype(str))
                label_encoders[feature] = le

            # 合并所有特征
            X = pd.concat([numeric_data, categorical_data], axis=1)

            # 创建目标变量：未来 3 个月资产是否达到 100 万+
            df['next_assets'] = df.groupby('customer_id')['total_assets'].shift(-prediction_period)
            df['target'] = (df['next_assets'] >= 1000000).astype(int)

            # 删除缺失值
            df.dropna(subset=['target'], inplace=True)
            X = X.loc[df.index]
            y = df['target']

            # 数据集划分
            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

            # 标准化数值特征
            scaler = StandardScaler()
            X_train_scaled = scaler.fit_transform(X_train)
            X_test_scaled = scaler.transform(X_test)

            # 训练 LightGBM 模型
            train_data = lgb.Dataset(X_train_scaled, label=y_train)
            test_data = lgb.Dataset(X_test_scaled, label=y_test, reference=train_data)

            params = {
                'objective': 'binary',
                'metric': 'auc',
                'boosting_type': 'gbdt',
                'learning_rate': 0.05,
                'num_leaves': 31,
                'feature_fraction': 0.9,
                'bagging_fraction': 0.8,
                'bagging_freq': 5,
                'verbose': -1
            }

            model = lgb.train(
                params,
                train_data,
                valid_sets=[test_data],
                num_boost_round=100,
                callbacks=[lgb.log_evaluation(0)]
            )

            # 特征重要性
            importance = model.feature_importance(importance_type='gain')
            feature_names = X.columns
            importance_df = pd.DataFrame({
                'feature': feature_names,
                'importance': importance
            }).sort_values(by='importance', ascending=False)

            # 可视化特征重要性
            plt.figure(figsize=(10, 6))
            plt.barh(range(len(importance_df)), importance_df['importance'], color='skyblue')
            plt.yticks(range(len(importance_df)), importance_df['feature'])
            plt.xlabel('Feature Importance (Gain)')
            plt.title('LightGBM Feature Importance for Predicting 1M+ Assets in {} Months'.format(prediction_period))
            plt.tight_layout()

            save_dir = os.path.join(os.path.dirname(__file__), 'image_show')
            os.makedirs(save_dir, exist_ok=True)
            filename = f'lightgbm_feature_importance_{prediction_period}.png'
            save_path = os.path.join(save_dir, filename)
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            plt.close()

            # 返回结果
            result = f"""
            LightGBM 分析完成！
            
            特征重要性排序：
            {importance_df.to_markdown(index=False)}
            
            特征重要性图表已保存到 {filename}
            """
            return result

        except Exception as e:
            return f"LightGBM 分析执行失败: {str(e)}"


# ====== 新增：SHAP 客户价值因素分析工具 ======
@register_tool('shap_analysis')
class SHAPAnalysisTool(BaseTool):
    """
    使用 SHAP 分析客户未来 3 个月资产提升至 100 万+的概率，并输出全局和局部解释。
    """
    description = '使用 SHAP 分析客户未来 3 个月资产提升至 100 万+的概率，并输出特征重要性'
    parameters = [{
        'name': 'prediction_period',
        'type': 'integer',
        'description': '预测周期（月数），默认为 3',
        'required': False
    }]

    def call(self, params: str, **kwargs) -> str:
        import json
        import lightgbm as lgb
        from sklearn.model_selection import train_test_split
        from sklearn.preprocessing import LabelEncoder, StandardScaler
        import matplotlib.pyplot as plt
        import shap

        args = json.loads(params)
        prediction_period = args.get('prediction_period', 3)  # 默认预测 3 个月

        # 从数据库获取客户数据
        engine = create_connection()
        if engine is None:
            return "数据库连接失败"

        # 获取客户历史数据
        query = f"""
        SELECT cb.customer_id, cb.age, cb.gender, cb.monthly_income, cb.occupation, 
               cba.total_assets, cba.deposit_balance, cba.financial_balance, 
               cba.fund_balance, cba.insurance_balance, cba.product_count,
               cba.financial_repurchase_count, cba.credit_card_monthly_expense,
               cba.investment_monthly_count, cba.app_login_count,
               cba.app_financial_view_time, cba.app_product_compare_count,
               cba.stat_month
        FROM customer_base cb
        JOIN customer_behavior_assets cba ON cb.customer_id = cba.customer_id
        ORDER BY cb.customer_id, cba.stat_month
        """

        try:
            df = pd.read_sql(query, engine)
            if df.empty:
                return "未找到任何客户数据"

            # 特征工程
            features_for_clustering = [
                'age', 'monthly_income', 'total_assets', 'deposit_balance', 
                'financial_balance', 'fund_balance', 'insurance_balance',
                'product_count', 'financial_repurchase_count', 
                'credit_card_monthly_expense', 'investment_monthly_count',
                'app_login_count', 'app_financial_view_time', 
                'app_product_compare_count'
            ]

            categorical_features = ['gender', 'occupation']

            # 处理数值特征
            numeric_data = df[features_for_clustering].fillna(0)

            # 处理分类特征 - 使用 LabelEncoder 编码
            categorical_data = df[categorical_features].copy()
            label_encoders = {}
            for feature in categorical_features:
                le = LabelEncoder()
                categorical_data[feature] = le.fit_transform(categorical_data[feature].astype(str))
                label_encoders[feature] = le

            # 合并所有特征
            X = pd.concat([numeric_data, categorical_data], axis=1)

            # 创建目标变量：未来 3 个月资产是否达到 100 万+
            df['next_assets'] = df.groupby('customer_id')['total_assets'].shift(-prediction_period)
            df['target'] = (df['next_assets'] >= 1000000).astype(int)

            # 删除缺失值
            df.dropna(subset=['target'], inplace=True)
            X = X.loc[df.index]
            y = df['target']

            # 数据集划分
            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

            # 标准化数值特征
            scaler = StandardScaler()
            X_train_scaled = scaler.fit_transform(X_train)
            X_test_scaled = scaler.transform(X_test)

            # 训练 LightGBM 模型
            train_data = lgb.Dataset(X_train_scaled, label=y_train)
            test_data = lgb.Dataset(X_test_scaled, label=y_test, reference=train_data)

            params = {
                'objective': 'binary',
                'metric': 'auc',
                'boosting_type': 'gbdt',
                'learning_rate': 0.05,
                'num_leaves': 31,
                'feature_fraction': 0.9,
                'bagging_fraction': 0.8,
                'bagging_freq': 5,
                'verbose': -1
            }

            model = lgb.train(
                params,
                train_data,
                valid_sets=[test_data],
                num_boost_round=100,
                callbacks=[lgb.log_evaluation(0)]
            )

            # 全局解释
            explainer = shap.TreeExplainer(model)
            shap_values = explainer.shap_values(X_train_scaled)

            # 可视化特征重要性
            plt.figure(figsize=(10, 6))
            shap.summary_plot(shap_values, X_train_scaled, feature_names=X.columns, show=False)
            plt.title('SHAP Feature Importance for Predicting 1M+ Assets in {} Months'.format(prediction_period))
            plt.tight_layout()

            save_dir = os.path.join(os.path.dirname(__file__), 'image_show')
            os.makedirs(save_dir, exist_ok=True)
            filename_global = f'shap_feature_importance_global_{prediction_period}.png'
            save_path_global = os.path.join(save_dir, filename_global)
            plt.savefig(save_path_global, dpi=300, bbox_inches='tight')
            plt.close()

            # 局部解释
            customer_index = 0
            customer_data = X_test.iloc[customer_index:customer_index+1]
            customer_shap_values = explainer.shap_values(customer_data)

            plt.figure(figsize=(12, 8))
            shap.waterfall_plot(shap.Explanation(values=customer_shap_values[0], 
                                                 base_values=explainer.expected_value, 
                                                 data=customer_data.iloc[0], 
                                                 feature_names=X.columns), show=False)
            plt.title(f'SHAP Local Explanation for Customer {customer_index}')
            plt.tight_layout()

            filename_local = f'shap_feature_importance_local_{customer_index}_{prediction_period}.png'
            save_path_local = os.path.join(save_dir, filename_local)
            plt.savefig(save_path_local, dpi=300, bbox_inches='tight')
            plt.close()

            # 返回结果
            result = f"""
            SHAP 分析完成！
            
            全局特征重要性图表已保存到 {filename_global}
            局部解释图表已保存到 {filename_local}
            """
            return result

        except Exception as e:
            return f"SHAP 分析执行失败: {str(e)}"

# ====== 初始化银行客户助手服务 ======
def init_agent_service():
    """初始化银行客户助手服务"""
    llm_cfg = {
        'model': 'qwen-plus-2025-04-28',
        'timeout': 30,
        'retry_count': 3,
    }
    tools = [
        {
            "mcpServers": {
                "tavily-mcp": {
                    # Tavily MCP配置
                    "command": "npx",
                    "args": ["-y", "tavily-mcp@0.1.4"],
                    "env": {
                        "TAVILY_API_KEY": os.getenv('TAVILY_API_KEY', ''),
                    },
                    "disabled": False,
                    "autoApprove": []
                },
            }
        },
        'exc_sql','lr_predict','lightgbm_analysis','shap_analysis'
    ]
    try:
        bot = Assistant(
            llm=llm_cfg,
            name='百万客群经营助手',
            description='银行客户数据查询与分析',
            system_message=system_prompt,
            function_list=tools,  # 确保工具已注册
           # files=['./faq.txt']
        )
        print("助手初始化成功！")
        return bot
    except Exception as e:
        print(f"助手初始化失败: {str(e)}")
        raise


# ====== 定义 app_gui 函数以启动 Web 界面 ======
def app_gui():
    """图形界面模式，提供 Web 图形界面"""
    try:
        print("正在启动 Web 界面...")
        # 初始化助手
        bot = init_agent_service()
        # 配置聊天界面，列举典型查询问题
        chatbot_config = {
            'prompt.suggestions': [
                '我行目前有多少客户？总资产管理规模是多少？',
                '客户的平均资产是多少？高净值客户的占比如何？',
                '客户年龄分布情况如何？哪个年龄段客户最多？',
                '客户持有的四类产品（存款、理财、基金、保险）分布情况如何？',
                '不同客户等级（普通、潜力、临界、高净值）的资产占比及客户数量是多少？',
                '影响客户价值的最重要因素有哪些？',
                '316d72c0795f4fa9a61aeb8804c21b8f这个客户未来三个月是否有可能资产达到100万？',
                '根据历史数据，未来3个月的客户资产总规模统计会如何变化？'
            ]
        }
        print("Web 界面准备就绪，正在启动服务...")
        # 启动 Web 界面
        WebUI(
            bot,
            chatbot_config=chatbot_config
        ).run()
    except Exception as e:
        print(f"启动 Web 界面失败: {str(e)}")
        print("请检查网络连接和 API Key 配置")


if __name__ == '__main__':
    # 运行模式选择
    app_gui()          # 图形界面模式（默认）